﻿// Gateway.Api 项目的版权、商标、专利和其他相关权利均受相应法律法规的保护。使用本项目应遵守相关法律法规和许可证的要求。
//
// 本项目主要遵循 MIT 许可证和 Apache 许可证（版本 2.0）进行分发和使用。许可证位于源代码树根目录中的 LICENSE-MIT 和 LICENSE-APACHE 文件。
//
// 不得利用本项目从事危害国家安全、扰乱社会秩序、侵犯他人合法权益等法律法规禁止的活动！任何基于本项目二次开发而产生的一切法律纠纷和责任，我们不承担任何责任！

using Aurora.AI.Domain.Domain.RateLimitModel.Aggregates;
using Aurora.AI.Domain.Domain.RateLimitModel.Repositories;
using Aurora.Redis;
using Microsoft.AspNetCore.Http;

namespace Aurora.AI.Gateway.Services;
/// <summary>
/// 模型限流服务
/// </summary>
public class RateLimitModelService(IRateLimitModelRepository repository, IRedisCache redisCache)
{
    private const string RatesCacheKey = "AI:GATEWAY:RATELIMIT";

    private async Task ReloadCache()
    {
        var data = await repository.GetAll();
        await redisCache.AddAsync(RatesCacheKey, data, TimeSpan.FromHours(24));
    }
    /// <summary>
    /// 获取速率限制
    /// </summary>
    /// <returns></returns>
    private async Task<List<RateLimitModel>> GetRatesByCache(bool reload = false)
    {
        var data = await redisCache.GetAsync<List<RateLimitModel>>(RatesCacheKey);
        if (data == null)
        {
            data = await repository.GetAll();
            await redisCache.AddAsync(RatesCacheKey, data, TimeSpan.FromHours(24));
        }
        return data;
    }
    /// <summary>
    /// 检查 模型速率
    /// </summary>
    /// <param name="model"></param>
    /// <param name="context"></param>
    /// <returns></returns>
    /// <exception cref="ArgumentException"></exception>
    /// <exception cref="Exception"></exception>
    public async ValueTask CheckAsync(string model, HttpContext context)
    {
        var ip = context.Connection.RemoteIpAddress?.ToString();
        if (context.Request.Headers.TryGetValue("X-Forwarded-For", out var header))
        {
            ip = header;
        }
        var limits = await GetRatesByCache();


        foreach (var item in limits.Where(x => x.Model.Contains(model)))
        {
            if (item.WhiteList.Contains(ip))
            {
                return;
            }
            if (item.BlackList.Contains(ip))
            {
                throw new ArgumentException("IP is in the blacklist");
            }
            string key = $"AI:GATEWAY:RATELIMIT_ITEM:{model}:{ip}{item.Id}";
            int? cache = await redisCache.GetAsync<int?>(key);
            if (cache.HasValue)
            {
                if (cache.Value >= item.Value)
                {
                    throw new ArgumentException("Rate limit exceeded");
                }
                await redisCache.IncrementAsync(key);
            }
            else
            {
                var strategy = item.Strategy switch
                {
                    "s" => TimeSpan.FromSeconds(item.Limit),
                    "m" => TimeSpan.FromMinutes(item.Limit),
                    "h" => TimeSpan.FromHours(item.Limit),
                    "d" => TimeSpan.FromDays(item.Limit),
                    _ => throw new Exception("Invalid strategy")
                };
                await redisCache.IncrementAsync(key, 1, strategy);
            }

        }


    }

    /// <summary>
    /// 获取分页列表
    /// </summary>
    /// <param name="page"></param>
    /// <param name="pageSize"></param>
    /// <returns></returns>
    public async ValueTask<(int total, List<RateLimitModel>)> GetPageListAsync(int page, int pageSize)
    {
        var data = await repository.GetPageListAsync(page, pageSize);
        return data;
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="rateLimitModel"></param>
    /// <returns></returns>
    public async ValueTask CreateAsync(RateLimitModel rateLimitModel)
    {
        await repository.Create(rateLimitModel);
        await ReloadCache();
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="id"></param>
    /// <returns></returns>
    public async ValueTask<bool> RemoveAsync(string id)
    {
        await repository.Remove(id);
        await ReloadCache();
        return true;
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="rateLimitModel"></param>
    /// <returns></returns>
    public async ValueTask<bool> UpdateAsync(RateLimitModel rateLimitModel)
    {
        await repository.UpdateLimit(rateLimitModel);
        await ReloadCache();
        return true;
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="id"></param>
    /// <returns></returns>
    public async ValueTask Disable(string id)
    {
        await repository.Disable(id);

        await ReloadCache();
    }

}
